import msgpack
from concurrent.futures import ThreadPoolExecutor
from collections import defaultdict
from threading import Thread
from typing import List, Optional

import zmq

from ctaf_core.ctaf_objects.base_event import BaseEvent, SysTerminateEvent, msgpack_event_default, \
    reconstruct_event_object_hook
from ctaf_core.ctaf_objects.ctaf_general_types import EventHandlerType
from ctaf_facade.enumerations.event_handler_exe_type_num import EventHandlerExeTypeEnum
from ctaf_facade.enumerations.event_type_enum import EventTypeEnum
from ctaf_core.log_trace_bundle.logging import LogRecorder, LoggerFactory



class EventEngine:
    """
    Event engine distributes event object based on its type
    to those handlers registered.
    """

    # 定义基于ZMQ的线程间inproc通信的地址
    _INPROC_EVENT_ADDR = "inproc://event_stream"

    log_recorder: LogRecorder = LoggerFactory.get_logger()

    def __init__(self, max_io_workers: int = 10, max_compute_workers: int = 5) -> None:

        self._active: bool = False
        self._event_dispatcher: Thread = Thread(target=self._run, name="EventDispatcher")

        self._handlers: defaultdict[EventTypeEnum, defaultdict[EventHandlerExeTypeEnum, List[EventHandlerType]]] = \
            defaultdict(lambda: defaultdict(list))
        self._general_handlers: defaultdict[EventHandlerExeTypeEnum, List[EventHandlerType]] = \
            defaultdict(list)

        self._io_executor: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=max_io_workers,
                                                                   thread_name_prefix="IOThreadPool")
        self._compute_executor: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=max_compute_workers,
                                                                        thread_name_prefix="ComputeThreadPool")

        # ZMQ Context 和 Sockets
        self._zmq_context: zmq.Context = zmq.Context()
        # Publisher socket for putting events into the internal stream
        # This will be created in put() or by external producers
        self._zmq_publisher: Optional[zmq.Socket] = None

        # Subscriber socket for the main event loop to receive events
        self._zmq_subscriber: Optional[zmq.Socket] = None

        # 用于停止ZMQ主循环的poller，允许更优雅地退出
        self._poller = zmq.Poller()
        self._stop_event = SysTerminateEvent()  # SysTerminateEvent now has to_dict/from_dict

    def _init_zmq_sockets(self):
        """Initializes ZMQ sockets. Called when the engine starts."""
        if self._zmq_publisher is None:  # Lazy initialization for publisher
            self._zmq_publisher = self._zmq_context.socket(zmq.PUB)
            # Bind the publisher to the inproc address
            self._zmq_publisher.bind(self._INPROC_EVENT_ADDR)
            EventEngine.log_recorder.write_info_log(f"EventEngine ZMQ publisher bind to {self._INPROC_EVENT_ADDR}")

        if self._zmq_subscriber is None:
            self._zmq_subscriber = self._zmq_context.socket(zmq.SUB)
            # Subscribe to all messages
            self._zmq_subscriber.setsockopt_string(zmq.SUBSCRIBE, "")
            # Connect the subscriber to the publisher's bind address
            self._zmq_subscriber.connect(self._INPROC_EVENT_ADDR)
            EventEngine.log_recorder.write_info_log(f"EventEngine ZMQ subscriber connected to {self._INPROC_EVENT_ADDR}")

            # 注册subscriber到poller，用于非阻塞等待事件
            self._poller.register(self._zmq_subscriber, zmq.POLLIN)

    def _run(self) -> None:
        """
        Get event from ZMQ inproc queue and then process it.
        """
        EventEngine.log_recorder.write_info_log("Event_dispatcher thread started.")
        self._init_zmq_sockets()  # 确保在主线程中初始化ZMQ sockets

        while self._active:
            try:
                # 使用poller进行非阻塞等待，timeout可以设为很小的值或None
                # 设为100ms可以周期性检查_active标志，同时避免CPU空转
                sockets = dict(self._poller.poll(timeout=100))

                if self._zmq_subscriber in sockets and sockets[self._zmq_subscriber] == zmq.POLLIN:
                    # 收到消息，接收并使用 msgpack 反序列化
                    msg_bytes = self._zmq_subscriber.recv()
                    # 使用 object_hook 来重建自定义对象
                    event: BaseEvent = msgpack.unpackb(msg_bytes, raw=False, object_hook=reconstruct_event_object_hook)

                    # 检查是否收到停止事件。
                    if isinstance(event, SysTerminateEvent):
                        EventEngine.log_recorder.write_info_log("Received stop event, initiating shutdown.")
                        self._active = False  # 设置active为False以退出循环
                        break  # 退出循环

                    self._process(event)
                else:
                    # 没有事件，继续循环检查_active
                    pass

            except zmq.ZMQError as ze:
                if ze.errno == zmq.ETERM:  # Context was terminated
                    EventEngine.log_recorder.write_warn_log("ZMQ context terminated during _run. Exiting loop.")
                    self._active = False
                else:
                    EventEngine.log_recorder.write_error_log(f"ZMQ error in event engine main loop: {ze}")
            except Exception as e:
                EventEngine.log_recorder.write_error_log(f"Error in event engine main loop: {e}")

        # 退出循环后清理ZMQ资源
        self._cleanup_zmq_sockets()
        EventEngine.log_recorder.write_info_log("Event_dispatcher thread stopped.")

    def _cleanup_zmq_sockets(self):
        """Cleans up ZMQ sockets and context."""
        if self._zmq_publisher:
            self._zmq_publisher.unbind(self._INPROC_EVENT_ADDR)  # 解除绑定
            self._zmq_publisher.close()
            self._zmq_publisher = None
            EventEngine.log_recorder.write_info_log("EventEngine ZMQ Publisher closed.")
        if self._zmq_subscriber:
            self._zmq_subscriber.disconnect(self._INPROC_EVENT_ADDR)  # 断开连接
            self._zmq_subscriber.close()
            self._zmq_subscriber = None
            EventEngine.log_recorder.write_info_log("EventEngine ZMQ Subscriber closed.")
        if self._zmq_context:
            self._zmq_context.term()  # 终止上下文
            self._zmq_context = None
            EventEngine.log_recorder.write_info_log("EventEngine ZMQ Context terminated.")

    def _process(self, event: BaseEvent) -> None:
        """
        Distribute event to those handlers registered based on their execution type.
        """
        event_type = event.get_event_type()

        if event_type in self._handlers:
            type_handlers = self._handlers[event_type]
            self._execute_handlers(event, type_handlers)

        if self._general_handlers:
            self._execute_handlers(event, self._general_handlers)

    def _execute_handlers(self, event: BaseEvent,
                          handlers_map: defaultdict[EventHandlerExeTypeEnum, List[EventHandlerType]]) -> None:
        """
        Execute handlers based on their execution type.
        """
        for handler in handlers_map[EventHandlerExeTypeEnum.SYNC_HANDLER]:
            try:
                handler(event)
            except Exception as e:
                EventEngine.log_recorder.write_error_log(
                    f"Sync handler {handler.__name__} for event {event.get_event_type()} failed: {e}")

        for handler in handlers_map[EventHandlerExeTypeEnum.ASYNC_IO_HANDLER]:
            try:
                self._io_executor.submit(self._safe_execute_handler, handler, event)
            except Exception as e:
                EventEngine.log_recorder.write_error_log(
                    f"Failed to submit async IO handler {handler.__name__} for event {event.get_event_type()}: {e}")

        for handler in handlers_map[EventHandlerExeTypeEnum.ASYNC_COMPUTE_HANDLER]:
            try:
                self._compute_executor.submit(self._safe_execute_handler, handler, event)
            except Exception as e:
                EventEngine.log_recorder.write_error_log(
                    f"Failed to submit async compute handler {handler.__name__} for event {event.get_event_type()}: {e}")

    def _safe_execute_handler(self, handler: EventHandlerType, event: BaseEvent) -> None:
        try:
            handler(event)
        except Exception as e:
            EventEngine.log_recorder.write_error_log(
                f"Async handler {handler.__name__} for event {event.get_event_type()} failed: {e}")

    def start(self) -> None:
        """
        Start event engine to process events.
        """
        if not self._active:
            self._active = True
            self._zmq_publisher = self._zmq_context.socket(zmq.PUB)
            self._zmq_publisher.bind(self._INPROC_EVENT_ADDR)  # PUB binds for inproc

            self._event_dispatcher.start()
            EventEngine.log_recorder.write_info_log("Event engine started.")

    def stop(self) -> None:
        """
        Stop event engine and shutdown thread pools.
        """
        if self._active:
            EventEngine.log_recorder.write_info_log("Initiating EventEngine stop.")
            # 1. Signal the event dispatcher to stop
            self._active = False
            if self._zmq_publisher:
                try:
                    # 使用 msgpack.dumps 序列化停止事件，并指定 default 编码器
                    self._zmq_publisher.send(msgpack.dumps(self._stop_event, default=msgpack_event_default))
                    EventEngine.log_recorder.write_info_log("Sent stop event via ZMQ.")
                except zmq.ZMQError as ze:
                    EventEngine.log_recorder.write_warn_log(f"ZMQ error sending stop event during shutdown: {ze}")
                except Exception as e:
                    EventEngine.log_recorder.write_error_log(f"Error sending stop event: {e}")

            # 2. Wait for the event dispatcher thread to finish its loop and clean up ZMQ
            if self._event_dispatcher.is_alive():
                self._event_dispatcher.join(timeout=5)  # Add a timeout in case it hangs
            if self._event_dispatcher.is_alive():
                EventEngine.log_recorder.write_warn_log(
                    "Event dispatcher thread did not terminate gracefully within timeout.")

            # 3. Shutdown thread pools
            EventEngine.log_recorder.write_info_log("Shutting down IO and Compute Thread Pools.")
            self._io_executor.shutdown(wait=True)
            self._compute_executor.shutdown(wait=True)

            # Ensure ZMQ context is terminated if _cleanup_zmq_sockets somehow missed it
            if self._zmq_context and not self._zmq_context.closed:
                EventEngine.log_recorder.write_info_log("Terminating ZMQ context as a fallback.")
                self._zmq_context.term()

            EventEngine.log_recorder.write_info_log("Event engine stopped and thread pools shut down.")

    def put(self, event: BaseEvent) -> None:
        """
        Put an event object into event queue via ZMQ.
        """
        if not self._active or not self._zmq_publisher:
            EventEngine.log_recorder.write_warn_log(
                "Event engine not active or ZMQ publisher not initialized. Event not put.")
            return

        try:
            # 序列化事件对象，使用 msgpack.dumps 并指定 default 编码器
            msg_bytes = msgpack.dumps(event, default=msgpack_event_default)
            self._zmq_publisher.send(msg_bytes)
            # EventEngine.log_recorder.write_debug_log(f"Event {event.get_event_type()} put to ZMQ.")
        except zmq.ZMQError as ze:
            EventEngine.log_recorder.write_error_log(f"Failed to put event into ZMQ queue: {ze}")
        except TypeError as te:  # Catch TypeError if _msgpack_event_default fails for an unhandled type
            EventEngine.log_recorder.write_error_log(
                f"Failed to serialize event with Msgpack: {te}. Event Type: {event.get_event_type()}")
        except Exception as e:
            EventEngine.log_recorder.write_error_log(f"Failed to put event: {e}")

    def register(self, event_type: EventTypeEnum,
                 event_handler: EventHandlerType,
                 execution_type: EventHandlerExeTypeEnum = EventHandlerExeTypeEnum.SYNC_HANDLER) -> None:
        handler_list: list = self._handlers[event_type][execution_type]
        if event_handler not in handler_list:
            handler_list.append(event_handler)
            EventEngine.log_recorder.write_debug_log(
                f"Registered handler {event_handler.__name__} for event type {event_type} as {execution_type}")

    def unregister(self, event_type: EventTypeEnum,
                   event_handler: EventHandlerType,
                   execution_type: EventHandlerExeTypeEnum = EventHandlerExeTypeEnum.SYNC_HANDLER) -> None:
        handler_list: list = self._handlers[event_type][execution_type]
        if event_handler in handler_list:
            handler_list.remove(event_handler)
            EventEngine.log_recorder.write_debug_log(
                f"Unregistered handler {event_handler.__name__} for event type {event_type} as {execution_type}")
            if not handler_list:
                self._handlers[event_type].pop(execution_type, None)
            if not self._handlers[event_type]:
                self._handlers.pop(event_type, None)

    def register_general(self, event_handler: EventHandlerType,
                         execution_type: EventHandlerExeTypeEnum = EventHandlerExeTypeEnum.SYNC_HANDLER) -> None:
        handler_list: list = self._general_handlers[execution_type]
        if event_handler not in handler_list:
            handler_list.append(event_handler)
            EventEngine.log_recorder.write_debug_log(
                f"Registered general handler {event_handler.__name__} as {execution_type}")

    def unregister_general(self, event_handler: EventHandlerType,
                           execution_type: EventHandlerExeTypeEnum = EventHandlerExeTypeEnum.SYNC_HANDLER) -> None:
        handler_list: list = self._general_handlers[execution_type]
        if event_handler in handler_list:
            handler_list.remove(event_handler)
            EventEngine.log_recorder.write_debug_log(
                f"Unregistered general handler {event_handler.__name__} as {execution_type}")
            if not handler_list:
                self._general_handlers.pop(execution_type, None)